from flask import Flask, request, send_from_directory, jsonify
import os
import threading
import asyncio
import client

app = Flask(__name__)

# 设置上传文件的保存目录
UPLOAD_FOLDER = '../data'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# 确保上传目录存在
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

# 文件上传接口
@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({"error": "没有文件部分"}), 400

    file = request.files['file']

    if file.filename == '':
        return jsonify({"error": "没有选择文件"}), 400

    if file:
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
        file.save(file_path)
        return jsonify({"message": "文件上传成功", "filename": file.filename}), 200

# 文件下载接口
@app.route('/download/<filename>', methods=['GET'])
def download_file(filename):
    file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    if not os.path.exists(file_path):
        return jsonify({"error": "文件不存在"}), 404
    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)

# 列出已上传文件接口
@app.route('/files', methods=['GET'])
def list_files():
    files = os.listdir(app.config['UPLOAD_FOLDER'])
    return jsonify({"files": files}), 200


@app.route("/send_prompt", methods=["POST"])
def send_prompt():
    data = request.json
    prompt = data.get('prompt', '')
    client.in_buf = prompt
    return "ok"


@app.route("/get_output", methods=["GET"])
def get_output():
    return jsonify({"output": client.out_buf}), 200


def run_flask():
    app.run(host='0.0.0.0', port=6678)


async def main():
    flask_thread = threading.Thread(target=run_flask, daemon=True)
    flask_thread.start()

    await client.main()


if __name__ == '__main__':
    asyncio.run(main())
    